#!/usr/bin/env python3
"""自定义异常和错误处理工具"""

import traceback
import sys
from typing import Optional, Dict, Any, Callable, Type
from functools import wraps
from pathlib import Path

from .logger import get_logger

logger = get_logger(__name__)


# 自定义异常类
class AutoVideoError(Exception):
    """基础异常类"""
    
    def __init__(self, message: str, error_code: Optional[str] = None, details: Optional[Dict[str, Any]] = None):
        super().__init__(message)
        self.message = message
        self.error_code = error_code or self.__class__.__name__
        self.details = details or {}
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            'error_type': self.__class__.__name__,
            'error_code': self.error_code,
            'message': self.message,
            'details': self.details
        }
    
    def __str__(self) -> str:
        if self.details:
            return f"{self.message} (Details: {self.details})"
        return self.message


class ModelError(AutoVideoError):
    """模型相关错误"""
    pass


class ModelLoadError(ModelError):
    """模型加载错误"""
    pass


class ModelNotFoundError(ModelError):
    """模型未找到错误"""
    pass


class GenerationError(AutoVideoError):
    """生成相关错误"""
    pass


class ImageGenerationError(GenerationError):
    """图像生成错误"""
    pass


class ConfigError(AutoVideoError):
    """配置相关错误"""
    pass


class ConfigLoadError(ConfigError):
    """配置加载错误"""
    pass


class ConfigValidationError(ConfigError):
    """配置验证错误"""
    pass


class FileError(AutoVideoError):
    """文件操作错误"""
    pass


class FileNotFoundError(FileError):
    """文件未找到错误"""
    pass


class FilePermissionError(FileError):
    """文件权限错误"""
    pass


class ValidationError(AutoVideoError):
    """验证错误"""
    pass


class ParameterError(ValidationError):
    """参数错误"""
    pass


class ResourceError(AutoVideoError):
    """资源相关错误"""
    pass


class MemoryError(ResourceError):
    """内存不足错误"""
    pass


class GPUError(ResourceError):
    """GPU相关错误"""
    pass


class NetworkError(AutoVideoError):
    """网络相关错误"""
    pass


class DownloadError(NetworkError):
    """下载错误"""
    pass


class ErrorHandler:
    """错误处理器"""
    
    def __init__(self):
        self.error_counts = {}
        self.error_handlers = {}
    
    def register_handler(self, error_type: Type[Exception], handler: Callable[[Exception], Any]) -> None:
        """注册错误处理器"""
        self.error_handlers[error_type] = handler
    
    def handle_error(self, error: Exception, context: Optional[str] = None) -> Any:
        """处理错误"""
        error_type = type(error)
        
        # 记录错误统计
        error_name = error_type.__name__
        self.error_counts[error_name] = self.error_counts.get(error_name, 0) + 1
        
        # 记录错误日志
        context_info = f" in {context}" if context else ""
        logger.error(f"{error_name}{context_info}: {error}")
        
        # 记录详细的错误信息
        if isinstance(error, AutoVideoError):
            logger.error(f"Error details: {error.details}")
        
        # 记录堆栈跟踪
        logger.debug(f"Stack trace:\n{traceback.format_exc()}")
        
        # 查找并执行错误处理器
        for registered_type, handler in self.error_handlers.items():
            if isinstance(error, registered_type):
                try:
                    return handler(error)
                except Exception as handler_error:
                    logger.error(f"Error handler failed: {handler_error}")
                    break
        
        # 如果没有找到处理器，重新抛出异常
        raise error
    
    def get_error_statistics(self) -> Dict[str, int]:
        """获取错误统计"""
        return self.error_counts.copy()
    
    def reset_statistics(self) -> None:
        """重置错误统计"""
        self.error_counts.clear()


class RetryManager:
    """重试管理器"""
    
    @staticmethod
    def retry_on_error(
        max_retries: int = 3,
        delay: float = 1.0,
        backoff_factor: float = 2.0,
        exceptions: tuple = (Exception,)
    ):
        """重试装饰器"""
        def decorator(func: Callable) -> Callable:
            @wraps(func)
            def wrapper(*args, **kwargs):
                last_exception = None
                current_delay = delay
                
                for attempt in range(max_retries + 1):
                    try:
                        return func(*args, **kwargs)
                    except exceptions as e:
                        last_exception = e
                        
                        if attempt < max_retries:
                            logger.warning(
                                f"Attempt {attempt + 1} failed for {func.__name__}: {e}. "
                                f"Retrying in {current_delay:.2f}s..."
                            )
                            import time
                            time.sleep(current_delay)
                            current_delay *= backoff_factor
                        else:
                            logger.error(
                                f"All {max_retries + 1} attempts failed for {func.__name__}"
                            )
                
                # 如果所有重试都失败，抛出最后一个异常
                raise last_exception
            
            return wrapper
        return decorator


def safe_execute(func: Callable, *args, default=None, log_errors: bool = True, **kwargs) -> Any:
    """安全执行函数"""
    try:
        return func(*args, **kwargs)
    except Exception as e:
        if log_errors:
            logger.error(f"Error executing {func.__name__}: {e}")
        return default


def validate_parameter(value: Any, param_name: str, expected_type: Type, allow_none: bool = False) -> None:
    """验证参数"""
    if value is None and allow_none:
        return
    
    if not isinstance(value, expected_type):
        raise ParameterError(
            f"Parameter '{param_name}' must be of type {expected_type.__name__}, "
            f"got {type(value).__name__}",
            details={'parameter': param_name, 'expected_type': expected_type.__name__, 'actual_type': type(value).__name__}
        )


def validate_file_path(file_path: Path, must_exist: bool = True, must_be_file: bool = True) -> None:
    """验证文件路径"""
    if not isinstance(file_path, (str, Path)):
        raise ParameterError(
            "File path must be a string or Path object",
            details={'actual_type': type(file_path).__name__}
        )
    
    path_obj = Path(file_path)
    
    if must_exist and not path_obj.exists():
        raise FileNotFoundError(
            f"File or directory does not exist: {file_path}",
            details={'path': str(file_path)}
        )
    
    if must_exist and must_be_file and not path_obj.is_file():
        raise FileError(
            f"Path is not a file: {file_path}",
            details={'path': str(file_path), 'is_dir': path_obj.is_dir()}
        )


def validate_image_size(width: int, height: int, min_size: int = 64, max_size: int = 2048) -> None:
    """验证图像尺寸"""
    validate_parameter(width, 'width', int)
    validate_parameter(height, 'height', int)
    
    if width < min_size or width > max_size:
        raise ValidationError(
            f"Width must be between {min_size} and {max_size}, got {width}",
            details={'width': width, 'min_size': min_size, 'max_size': max_size}
        )
    
    if height < min_size or height > max_size:
        raise ValidationError(
            f"Height must be between {min_size} and {max_size}, got {height}",
            details={'height': height, 'min_size': min_size, 'max_size': max_size}
        )
    
    # 检查是否为8的倍数（Stable Diffusion要求）
    if width % 8 != 0:
        raise ValidationError(
            f"Width must be divisible by 8, got {width}",
            details={'width': width}
        )
    
    if height % 8 != 0:
        raise ValidationError(
            f"Height must be divisible by 8, got {height}",
            details={'height': height}
        )


def validate_prompt(prompt: str, max_length: int = 1000) -> None:
    """验证提示词"""
    validate_parameter(prompt, 'prompt', str)
    
    if not prompt.strip():
        raise ValidationError(
            "Prompt cannot be empty",
            details={'prompt': prompt}
        )
    
    if len(prompt) > max_length:
        raise ValidationError(
            f"Prompt too long (max {max_length} characters), got {len(prompt)}",
            details={'prompt_length': len(prompt), 'max_length': max_length}
        )


def handle_gpu_error(func: Callable) -> Callable:
    """GPU错误处理装饰器"""
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except RuntimeError as e:
            error_msg = str(e).lower()
            
            if 'out of memory' in error_msg or 'cuda out of memory' in error_msg:
                raise MemoryError(
                    "GPU out of memory. Try reducing batch size or image resolution.",
                    details={'original_error': str(e)}
                )
            elif 'cuda' in error_msg:
                raise GPUError(
                    f"CUDA error occurred: {e}",
                    details={'original_error': str(e)}
                )
            else:
                # 重新抛出原始错误
                raise
    
    return wrapper


def handle_model_error(func: Callable) -> Callable:
    """模型错误处理装饰器"""
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            error_msg = str(e).lower()
            
            if 'not found' in error_msg or 'no such file' in error_msg:
                raise ModelNotFoundError(
                    f"Model not found: {e}",
                    details={'original_error': str(e)}
                )
            elif 'permission' in error_msg or 'access' in error_msg:
                raise FilePermissionError(
                    f"Permission denied when accessing model: {e}",
                    details={'original_error': str(e)}
                )
            elif 'connection' in error_msg or 'network' in error_msg:
                raise NetworkError(
                    f"Network error when loading model: {e}",
                    details={'original_error': str(e)}
                )
            else:
                raise ModelLoadError(
                    f"Failed to load model: {e}",
                    details={'original_error': str(e)}
                )
    
    return wrapper


def create_error_context(context_name: str):
    """创建错误上下文装饰器"""
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args, **kwargs):
            try:
                return func(*args, **kwargs)
            except Exception as e:
                # 添加上下文信息到异常
                if isinstance(e, AutoVideoError):
                    e.details['context'] = context_name
                
                logger.error(f"Error in {context_name}: {e}")
                raise
        
        return wrapper
    return decorator


# 全局错误处理器实例
_global_error_handler = ErrorHandler()


def get_global_error_handler() -> ErrorHandler:
    """获取全局错误处理器"""
    return _global_error_handler


def setup_default_error_handlers() -> None:
    """设置默认错误处理器"""
    handler = get_global_error_handler()
    
    # GPU内存不足处理器
    def handle_gpu_memory_error(error: MemoryError) -> None:
        logger.warning("GPU memory error detected, clearing cache...")
        try:
            import torch
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                torch.cuda.synchronize()
        except ImportError:
            pass
    
    # 模型加载错误处理器
    def handle_model_load_error(error: ModelLoadError) -> None:
        logger.warning("Model load error detected, suggesting fallback options...")
        logger.info("Try using a different model or check your internet connection")
    
    # 注册处理器
    handler.register_handler(MemoryError, handle_gpu_memory_error)
    handler.register_handler(ModelLoadError, handle_model_load_error)


# 异常处理上下文管理器
class exception_handler:
    """异常处理上下文管理器"""
    
    def __init__(self, 
                 context: str,
                 reraise: bool = True,
                 log_errors: bool = True,
                 default_return=None):
        self.context = context
        self.reraise = reraise
        self.log_errors = log_errors
        self.default_return = default_return
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            if self.log_errors:
                logger.error(f"Exception in {self.context}: {exc_val}")
                logger.debug(f"Stack trace:\n{traceback.format_exc()}")
            
            if not self.reraise:
                return True  # 抑制异常
        
        return False  # 不抑制异常


# 初始化默认错误处理器
setup_default_error_handlers()